package com.zjh.admin.websocket.inteceptor;

import lombok.extern.slf4j.Slf4j;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.util.MultiValueMap;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.server.HandshakeInterceptor;
import org.springframework.web.util.UriComponents;
import org.springframework.web.util.UriComponentsBuilder;

import java.util.Map;

@Slf4j
public class CustomHandshakeInterceptor implements HandshakeInterceptor {

    @Override
    public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response,
                                   WebSocketHandler wsHandler, Map<String, Object> attributes) throws Exception {

        // 解析查询参数
        UriComponents uriComponents = UriComponentsBuilder.fromUri(request.getURI()).build();
        MultiValueMap<String, String> queryParams = uriComponents.getQueryParams();

        // 获取username（兼容旧版）
        String username = queryParams.getFirst("username");
        if (username != null) {
            attributes.put("username", username);
        }

        if (request.getURI().getQuery() != null) {
            attributes.put("ipAddress", request.getLocalAddress());
        }

        // 获取docId（新版功能）
        String docId = queryParams.getFirst("docId");

        if (docId != null) {
            attributes.put("docId", docId);
        }

        return true; // 返回false可阻止连接

    }

    @Override
    public void afterHandshake(ServerHttpRequest request, ServerHttpResponse response,
                               WebSocketHandler wsHandler, Exception exception) {
        // 可以在此处执行握手完成后的操作
    }
}